"""
人脸特征提取 UI 模块
"""

import streamlit as st
from PIL import Image
import numpy as np
from typing import Optional

from .api_client import APIClient
from .utils import image_to_base64, file_to_base64, display_json_result
import base64
import io


class FaceFeatureUI:
    """人脸特征提取 UI 组件"""

    def __init__(self, api_client: APIClient):
        """
        初始化人脸特征提取 UI

        Args:
            api_client: API 客户端实例
        """
        self.api_client = api_client

    def render(self):
        """渲染人脸特征提取 UI"""
        st.markdown(
            '<div class="main-header">🎭 人脸特征提取</div>', unsafe_allow_html=True
        )

        with st.expander("ℹ️ 功能说明", expanded=True):
            st.markdown("""
            上传包含人脸的图片，提取人脸特征向量。支持多种检测和特征提取模型。
            
            **使用步骤**：
            1. 上传图片
            2. 设置参数
            3. 点击"提取特征"按钮
            4. 查看结果
            """)

        # 图片上传区域
        col1, col2 = st.columns([1, 1])

        with col1:
            uploaded_file = st.file_uploader(
                "上传图片", type=["jpg", "jpeg", "png"], key="face_feature_upload"
            )

            if uploaded_file is not None:
                try:
                    image = Image.open(uploaded_file)
                    with col2:
                        st.image(image, caption="上传的图片", use_container_width=True)

                    # 参数设置
                    with st.expander("⚙️ 参数设置", expanded=True):
                        detection_model = st.selectbox(
                            "检测模型",
                            ["retinaface", "mtcnn", "dlib"],
                            index=0,
                            help="选择人脸检测使用的模型",
                        )
                        extract_model = st.selectbox(
                            "特征提取模型",
                            ["arcface", "facenet", "vggface"],
                            index=0,
                            help="选择特征提取使用的模型",
                        )
                        min_confidence = st.slider(
                            "最小置信度",
                            min_value=0.0,
                            max_value=1.0,
                            value=0.5,
                            step=0.05,
                            help="人脸检测的最小置信度阈值",
                        )

                    # 执行按钮
                    if st.button("🚀 提取特征", key="extract_feature_btn"):
                        with st.spinner("🔍 正在提取人脸特征..."):
                            # 将图片转换为 base64
                            image_base64 = file_to_base64(uploaded_file)

                            # 准备参数
                            params = {
                                "detection_model": detection_model,
                                "extract_model": extract_model,
                                "min_confidence": min_confidence,
                            }

                            # 调用 API
                            result = self.api_client.extract_face_features(
                                image_base64, params
                            )

                            # 显示结果
                            data = display_json_result(result, "✅ 特征提取完成！")

                            # 如果成功，显示更多信息
                            if data:
                                st.markdown("### 提取结果")

                                # 显示检测到的人脸
                                if "face_image" in data:
                                    try:
                                        face_image_data = base64.b64decode(
                                            data["face_image"]
                                        )
                                        face_image = Image.open(
                                            io.BytesIO(face_image_data)
                                        )
                                        st.image(
                                            face_image,
                                            caption="检测到的人脸",
                                            width=200,
                                        )
                                    except Exception as e:
                                        st.warning(f"无法显示人脸图像: {str(e)}")

                                # 显示特征向量
                                if "feature_vector" in data:
                                    st.markdown("#### 特征向量")
                                    feature_vector = data["feature_vector"]

                                    # 显示特征向量的前10个值
                                    st.write(f"维度: {len(feature_vector)}")
                                    st.write(f"前10个值: {feature_vector[:10]}")

                                    # 可视化特征向量
                                    st.line_chart(feature_vector[:50])

                except Exception as e:
                    st.error(f"处理图片时出错: {str(e)}")
            else:
                with col2:
                    st.info("请上传一张包含人脸的图片")
